import random

from pymilvus import (
    connections,
    utility,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
)


def get_test_collection() -> Collection:
    print('=== start connecting to Milvus     ===')
    # The maximum number of connections is 65,536
    connections.connect("default", host="ty.bmzhjt.cn", port="19530")

    fields = [
        FieldSchema(name="pk", dtype=DataType.INT64, is_primary=True, auto_id=False),
        FieldSchema(name="random", dtype=DataType.DOUBLE),
        FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=8)
    ]
    schema = CollectionSchema(fields, "hello_milvus is the simplest demo to introduce the APIs")
    return Collection("hello_milvus", schema)


hello_milvus = get_test_collection()

entities = [
    [i for i in range(3000)],  # field pk
    [float(random.randrange(-20, -10)) for _ in range(3000)],  # field random
    [[random.random() for _ in range(8)] for _ in range(3000)],  # field embeddings
]


def add_data():
    print('=== Start inserting entities       ===')
    insert_result = hello_milvus.insert(entities)
    hello_milvus.flush()


def create_index():
    print('=== Start Creating index IVF_FLAT  ===')
    index = {
        "index_type": "IVF_FLAT",
        "metric_type": "L2",
        "params": {"nlist": 128},
    }
    hello_milvus.create_index("embeddings", index)


def vectors_search():
    print('=== Start searching based on vector similarity ===')
    vectors_to_search = entities[-1][-2:]
    search_params = {
        "metric_type": "L2",
        "params": {"nprobe": 10},
    }
    result = hello_milvus.search(vectors_to_search, "embeddings", search_params, limit=3, output_fields=["random"])
    print(result)


def field_search():
    print('=== Start querying with `random > -14` ===')
    result = hello_milvus.query(expr="random > -14", output_fields=["random"])
    print(result)


def hybrid_search():
    print('=== Start hybrid searching with `random > -12` ===')
    vectors_to_search = entities[-1][-2:]
    search_params = {
        "metric_type": "L2",
        "params": {"nprobe": 10},
    }
    result = hello_milvus.search(vectors_to_search, "embeddings", search_params, limit=3, expr="random > -12",
                                 output_fields=["random"])
    print(result)


def delete_data():
    print('=== Start deleting with expr `pk in [0, 1]` ===')
    expr = f"pk in [{entities[0][0]}, {entities[0][1]}]"
    hello_milvus.delete(expr)


def drop_hello():
    print('=== Drop collection `hello_milvus` ===')
    utility.drop_collection("hello_milvus")


add_data()
create_index()

# 查询之前必须要load
print('=== Start loading collection `hello_milvus`  ===')
hello_milvus.load()

vectors_search()
field_search()
hybrid_search()

print('=== release collection `hello_milvus`  ===')
hello_milvus.release()

delete_data()
drop_hello()

print('=== disconnect  connection ===')
connections.disconnect("default")
